import sys
import os
sys.path.append(os.path.realpath(__file__))
import argparse

import random

from utils import train_util
import torch.distributed as dist

import torch.utils.data
import torch.utils.data.distributed
from torchvision import transforms

import yaml
from datasets.s3dis_closer_logger import setup_logger

from datasets.s3dis_closer_utils import PointcloudToTensor
from datasets.s3dis_closer import S3DISSeg
from datasets.s3dis_closer_train import train, validate, MaskedCrossEntropy

import numpy as np

torch.set_num_threads(1)


parser = argparse.ArgumentParser()
parser.add_argument("-c", "--config", help="path to py model", required=True, default='.configs/config_seg.yaml')
parser.add_argument("exp_name", help="name of the exp")

args = parser.parse_args()
config_path = args.config

assert(torch.cuda.device_count() == 1)
torch.cuda.set_device(0)


# torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
np.random.seed(42)
torch.cuda.manual_seed_all(42)
torch.manual_seed(42)
random.seed(42)
os.environ['PYTHONHASHSEED'] = str(42)

with open(config_path, 'r') as ymlfile:
    cfg = yaml.load(ymlfile, Loader=yaml.FullLoader)

# check if we masked all other devices

gen_model_file, gen_model_name = train_util.check_model_paths(cfg['model']['generator'])[0]


exp_descr = args.exp_name

aug = False
if 'aug' in cfg['data']:
    aug = cfg['data']['aug']

print('aug', aug)


class FakeCFG:
    def __init__(self):
        self.data_root = cfg['data']['path']
        self.batch_size = cfg['data']['batch_size']
        self.num_points = cfg['data']['num_points']
        self.epochs = cfg['train']['num_epochs']
        self.num_workers = cfg['data']['num_workers']

        self.num_steps = 2000
        self.input_features_dim = 4

        self.num_classes = 13

        self.sampleDl = 0.04

        self.in_radius = 2.0

        self.print_freq = 10

        self.x_angle_range = 0.0
        self.y_angle_range = 0.0
        self.z_angle_range = 3.1415926

        self.scale_low = 0.7
        self.scale_high = 1.3

        self.noise_std = 0.001
        self.noise_clip = 0.05
        self.translate_range = 0.0

        self.color_drop = 0.2

        self.augment_symmetries = [1, 0, 0]



config = FakeCFG()

test_transforms = transforms.Compose([
    PointcloudToTensor(),
])

val_dataset = S3DISSeg(input_features_dim=config.input_features_dim,
                       subsampling_parameter=config.sampleDl, color_drop=config.color_drop,
                       in_radius=config.in_radius, num_points=config.num_points,
                       num_steps=config.num_steps, num_epochs=20,
                       data_root=config.data_root, transforms=test_transforms,
                       split='val')

# val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)
val_loader = torch.utils.data.DataLoader(val_dataset,
                                         batch_size=config.batch_size,
                                         shuffle=False,
                                         num_workers=config.num_workers,
                                         pin_memory=True,
                                         # sampler=val_sampler,
                                         drop_last=False)
# model part
del cfg['model']['generator']
params_dict_gen = cfg['model']


writer, exp_dir, full_desc = train_util.create_experiment(exp_descr,
                                                          params_dict={'exp_root': cfg['experiment']['root'],
                                                                       'writer_root': cfg['experiment']['writer_root'],
                                                                       'config_path': config_path})
generator = train_util.get_model(gen_model_file, params_dict_gen,
                                 exp_dir=exp_dir).cuda()

if 'restore' in cfg:
    print('restoring')
    g_ckpt_path = cfg['restore']['generator']
    train_util.restore_exp_fix(objects=[generator],
                               names=[g_ckpt_path])

print('generator_params', sum(p.numel() for p in generator.parameters()))

runing_vote_logits = [np.zeros((config.num_classes, l.shape[0]), dtype=np.float32) for l in
                      val_loader.dataset.sub_clouds_points_labels]


show_each = cfg['train']['show_each']
show_iters = 0
data_iters = 0


label_smooth = False

if 'label_smooth' in cfg['train']:
    label_smooth = cfg['train']['label_smooth']


criterion = MaskedCrossEntropy()
model = generator

logger = setup_logger(output=None, distributed_rank=0, name="s3dis")

validate('Last', val_loader, model, criterion, runing_vote_logits, config,
         logger=logger, num_votes=20, save_path=exp_dir)
